package com.zoe.onelink.logging.trace.http.util;

import brave.Span;
import brave.Tracer;
import brave.internal.Platform;
import brave.internal.codec.HexCodec;
import brave.propagation.TraceContext;
import com.alibaba.fastjson.JSON;
import com.zoe.onelink.common.constant.TraceSpanConstant;
import com.zoe.onelink.core.util.http.HttpClient;
import com.zoe.onelink.core.util.http.HttpClientAsyncCallback;
import com.zoe.onelink.core.util.http.model.HttpClientEntity;
import com.zoe.onelink.core.util.http.model.HttpClientResponse;
import com.zoe.onelink.logging.trace.property.TraceLogProperties;
import com.zoe.onelink.util.StrUtil;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.SmartInitializingSingleton;

import java.util.Map;

/**
 * <p>标题: HTTP请求组件(带链路日志)</p>
 * <p>描述: </p>
 * <p>版权: Copyright (c) 2021</p>
 * <p>公司: 智业软件股份有限公司</p>
 *
 * @version: 1.0
 * @author: ruansheng
 * @date: 2021-10-18
 */
@Slf4j
@RequiredArgsConstructor
public class TracedOnelinkHttpClient implements HttpClient, SmartInitializingSingleton {


    /**
     * 链路追踪对象
     */
    @Getter
    private final HttpClient delegateClient;
    private final BeanFactory beanFactory;

    private Tracer tracer;
    private TraceLogProperties traceLogProperties;


    @Override
    public <R> HttpClientResponse<R> doRequest(HttpClientEntity httpClientEntity, Class<R> returnType) {
        Span span = this.nextSpan(httpClientEntity, false);
        try {
            HttpClientResponse<R> result = this.delegateClient.doRequest(httpClientEntity, returnType);
            this.logResult(span, result.getData());
            return result;
        } catch (Exception e) {
            span.error(e);
            throw e;
        } finally {
            span.finish();
        }
    }

    @Override
    public <R> void doAsyncRequest(HttpClientEntity httpClientEntity, Class<R> returnType, HttpClientAsyncCallback<R> callback) {
        final Span span = this.nextSpan(httpClientEntity, true);
        this.delegateClient.doAsyncRequest(httpClientEntity, returnType, new TracedAsyncAsyncCallback<>(span, callback));
    }

    /**
     * 记录请求返回结果
     */
    private <R> void logResult(Span span, R result) {
        if (result == null) {
            return;
        }
        String resultJsonStr = result instanceof String ? result.toString() : JSON.toJSONString(result);
        int resultMaxLength = this.traceLogProperties.getHttp().getResultMaxLength();
        if (resultJsonStr.length() > resultMaxLength) {
            span.tag(TraceSpanConstant.HTTP_RESULT, resultJsonStr.substring(0, resultMaxLength));
        } else {
            span.tag(TraceSpanConstant.HTTP_RESULT, resultJsonStr);
        }
    }

    /**
     * 生成下一个链路节点
     */
    private Span nextSpan(HttpClientEntity httpClientEntity, boolean isAsync) {
        // 创建新的Span节点
        Span span = this.createNextTraceSpan(httpClientEntity, isAsync);
        // 设置请求头
        this.setRequestHeader(span, httpClientEntity);
        // 设置链路信息
        this.setTraceInfo(span, httpClientEntity);

        return span;
    }

    /**
     * 设置链路信息
     *
     * @param span             链路节点
     * @param httpClientEntity 请求参数实体
     */
    private void setTraceInfo(Span span, HttpClientEntity httpClientEntity) {
        // 记录请求URL
        if (httpClientEntity.getUrl() != null) {
            span.tag(TraceSpanConstant.HTTP_URL, httpClientEntity.getUrl());
        }
        // 记录请求方式
        if (httpClientEntity.getMethod() != null) {
            span.tag(TraceSpanConstant.HTTP_METHOD, httpClientEntity.getMethod().name());
        }
        // 记录请求头
        Map<String, String> headerMap = httpClientEntity.getHeaderMap();
        if (headerMap != null && headerMap.size() != 0) {
            span.tag(TraceSpanConstant.HTTP_HEADER, JSON.toJSONString(headerMap));
        }
        // 记录请求参数
        if (httpClientEntity.getBody() != null) {
            span.tag(TraceSpanConstant.HTTP_PARAM, httpClientEntity.getBody());
        } else {
            span.tag(TraceSpanConstant.HTTP_PARAM, JSON.toJSONString(httpClientEntity.getFormMap()));
        }
    }

    /**
     * 设置链路信息
     */
    private void setRequestHeader(Span span, HttpClientEntity httpClientEntity) {
        Map<String, String> headerMap = httpClientEntity.getHeaderMap();
        if (headerMap == null) {
            return;
        }
        try {
            TraceContext context = span.context();
            // 设置链路ID
            headerMap.put(TraceSpanConstant.B3Propagation.TRACE_ID, context.traceIdString());
            // 为新的连接设置新的SpanId
            headerMap.put(TraceSpanConstant.B3Propagation.SPAN_ID, HexCodec.toLowerHex(Platform.get().randomLong()));
            // 将当前SpanId设置到下一个请求的ParentSpanId
            headerMap.put(TraceSpanConstant.B3Propagation.PARENT_SPAN_ID, context.spanIdString());
        } catch (Exception e) {
            log.warn("[HTTP-CLIENT] 设置链路日志请求头失败", e);
        }

    }


    /**
     * 创建下一个Zipkin链路Span
     */
    private Span createNextTraceSpan(HttpClientEntity httpClientEntity, boolean isAsync) {
        String requestUrl = httpClientEntity.getUrl().length() > 64 ? httpClientEntity.getUrl().substring(0, 64) : httpClientEntity.getUrl();
        String spanPrefix = isAsync ? "HTTP-CLIENT-ASYNC" : "HTTP-CLIENT";
        String spanName = StrUtil.format("[{}] {}", spanPrefix, requestUrl);
        return this.tracer.nextSpan()
                .name(spanName)
                .kind(Span.Kind.CLIENT)
                .tag(TraceSpanConstant.LOG_TYPE, TraceSpanConstant.LogType.THIRD_PARTY)
                .start();
    }

    @Override
    public void afterSingletonsInstantiated() {
        this.tracer = this.beanFactory.getBean(Tracer.class);
        this.traceLogProperties = this.beanFactory.getBean(TraceLogProperties.class);
    }

    public static class TracedAsyncAsyncCallback<R> implements HttpClientAsyncCallback<R> {

        private final Span span;
        private final HttpClientAsyncCallback<R> callback;

        public TracedAsyncAsyncCallback(Span span, HttpClientAsyncCallback<R> callback) {
            this.span = span;
            this.callback = callback;
        }

        @Override
        public void onSuccess(HttpClientResponse<R> response) {
            this.span.tag(TraceSpanConstant.HTTP_RESULT, StrUtil.nullToEmpty(JSON.toJSONString(response.getData())));
            try {
                if (callback != null) {
                    callback.onSuccess(response);
                }
            } catch (Exception e) {
                this.span.error(e);
            } finally {
                this.span.finish();
            }
        }

        @Override
        public void onFailure(Throwable exception) {
            this.span.error(exception);
            try {
                if (this.callback != null) {
                    this.callback.onFailure(exception);
                }
            } finally {
                this.span.finish();
            }
        }
    }


}
